import math
import threading
import time

import rtde_receive
import rtde_control
from Hardware import FT300
from Hardware.BaseFtSensor import BaseFt
import numpy as np
from rtde_control import RTDEControlInterface as RTDEControl
from Algorithm.AlgorithmModel import ControlResult
from Algorithm.AlgorithmModel import ControlMode
from Common.MatrixProcess import trans_pose
import Common.MatrixProcess as mp

class UrRobot():
    def __init__(self, robotip="192.168.68.250", forcesensor: BaseFt = None,dt=0.008, gripper=None, controlfuc=None):
        self.maxVel = 0.5
        self.maxAcc = 0.1
        self.maxAdjustPos = 0.005#移动
        self.maxAdjustAngle=0.01
        self.robotip = robotip
        self.ft = forcesensor
        self._controlfun=controlfuc

        self._startServoControl=False
        self._controlResult=ControlResult.Stop
        self.isSave=False
        self._logData=[]
        self.dt=dt
    def init(self):
        self.connect()
        self.FtInit()
        self.FtZero()
        time.sleep(0.5)

    def connect(self):
        self.rtde_r = rtde_receive.RTDEReceiveInterface(self.robotip)
        self.rtde_c = rtde_control.RTDEControlInterface(self.robotip)
        self.ft.SetRobot(self.rtde_c,self.rtde_r)
        self._startReadLoop = True
        self._readThread = threading.Thread(target=self._readLoop)
        self._readThread.start()
        print(123)

    def disconnect(self):
        self._startReadLoop=False
        self.rtde_r.disconnect()
        self.rtde_c.disconnect()
        # self.rtde_c.stopScript()
        self.ft.DisConnect()

        ####Motion####

    #region 控制类代码

    def initPeriod(self):
        return self.rtde_c.initPeriod()

    def waitPeriod(self, t_start):
        self.rtde_c.waitPeriod(t_start)

    def servoStop(self):
        self.rtde_c.servoStop()

    def moveJ(self, jointpos, vel, acc, asynchronous=False):
        vel = vel if vel <= self.maxVel else self.maxVel
        acc = acc if acc <= self.maxAcc else self.maxAcc
        self.rtde_c.moveJ(jointpos, vel, acc, asynchronous)

    def moveL(self,pose,vel,acc,asynchronous=False):
        vel = vel if vel <= self.maxVel else self.maxVel
        acc = acc if acc <= self.maxAcc else self.maxAcc
        self.rtde_c.moveL(pose, vel, acc, asynchronous)

    def servoL(self, poses, speed=0, acc=0, time=0.002, lookahead_time=0.1, gain=100):
        # poses = poses if poses <= self.maxAdjust else self.maxAdjust
        self.rtde_c.servoL(poses, speed, acc, time, lookahead_time, gain)

    def servoLAT(self,curpose,adjpose,speed,acc,time,lookahead_time,gain):
        # 平移不要大于0.005，旋转不要大于0.01
        adj=self.__check_ServoL_Adj(adjpose)
        #pose=self.rtde_c.poseTrans(curpose,adj)#基于tcp运动，右乘
        pose=trans_pose(curpose,adj)
        #print(pose)
        self.servoL(pose,speed,acc,time,lookahead_time,gain)

    def servoLAB(self,curpose,adjpose,speed,acc,time,lookahead_time,gain):
        '''
        以基座标的方向移动
        @param curpose: 当前tcp位置，通过gettcppose获取
        @param adjpose: 调整的增量
        @param speed: 不用，输入0
        @param acc: 不用，输入0
        @param time: 机器人的控制周期，100hz
        @param lookahead_time: 越小速度越快，太小会超调
        @param gain: 越大速度越快
        @return:
        '''
        pose=self.__check_ServoL_Adj(adjpose)
        #goalpose = self.rtde_c.poseTrans(pose, curpose)#基于基坐标系运动，左乘
        goalpose=pose+curpose
        self.servoL(goalpose, speed, acc, time, lookahead_time, gain)

    def servoJ(self, joints, speed, acc, time, lookahead_time, gain):
        joints = joints if joints <= self.maxAdjust else self.maxAdjust
        self.rtde_c.servoL(joints, speed, acc, time, lookahead_time, gain)
#endregion

    #region 读取类的代码

    def getActualQ(self, isRad=True) -> np.array:
        if (isRad == True):
            return np.array(self._actualQ)
        else:
            return np.array(self._actualQ) * 180 / math.pi

    def getActualQd(self):
        return np.array(self._actualQd)

    def getTargetQd(self):
        # Target joint velocities
        return np.array(self._targetQd)

    def getTargetQdd(self):
        return np.array(self._targetQdd)

    def getActualTCPPose(self):
        return np.array(self._actualTcpPose)

    def getActualTCPSpeed(self):
        return np.array(self._actualTcpSpeed)

    ####ForceSensor####
    def FtInit(self):
        pass

    def FtZero(self):
        self.ft.FtZero()

    def GetForce(self):
        return self.ft.GetForce()

    def GetForceBase(self):
        ft=self.ft.GetForce().copy()
        pose=self._actualTcpPose.copy()
        RBase_Tcp = mp.rotation_matrix_from_vectorr(pose[3:])  # 旋转向量转旋转矩阵
        F = RBase_Tcp.dot(ft[:3])  # 转化力矩
        T = RBase_Tcp.dot(ft[3:])  # 转化力矩
        ft[:3] = F
        ft[3:] = T
        return ft.copy()
#endregion

    #region 业务逻辑处理代码
    def __check_ServoL_Adj(self,adj):
        '''
        1，对调整量进行限制，幅值不能大于设定值
        2，需要保证移动的总方向不变
        @param adj:调整量
        @return:
        '''
        pose = np.array(adj)
        p = pose[:3]
        a = pose[-3:]
        sp = np.sign(p)
        sa = np.sign(a)
        #计算移动的幅值
        sump=np.linalg.norm(p)
        if(sump>self.maxAdjustPos):
            p=(self.maxAdjustPos/sump)*p
        #计算角度的幅值
        suma=np.linalg.norm(a)
        if(suma>self.maxAdjustAngle):
            a=(self.maxAdjustAngle/suma)*a
        # p[abs(p) > self.maxAdjustPos] = (self.maxAdjustPos * sp)
        # a[abs(a) > self.maxAdjustAngle] = (self.maxAdjustAngle * sa)
        pose = np.concatenate((p, a))
        return pose

    def _readLoop(self):
        while(self._startReadLoop==True):
            start = time.perf_counter()
            time0=self.initPeriod()
            self._actualQ=self.rtde_r.getActualQ()
            self._actualQd=self.rtde_r.getActualQd()
            self._targetQd=self.rtde_r.getTargetQd()
            self._targetQdd=self.rtde_r.getTargetQdd()
            self._actualTcpPose=self.rtde_r.getActualTCPPose()
            self._actualTcpSpeed=self.rtde_r.getActualTCPSpeed()
            self._ft=self.ft.GetForce()
            if(self._startServoControl==True):
                self._realTimeControl()

            if(self.isSave==True):
                logdata={'actualTcpPose':self._actualTcpPose,
                         'actualTcpSpeed':self._actualTcpSpeed,
                         'force':self._ft.copy()}
                self._logData.append(logdata)

            self.waitPeriod(time0)
            end = time.perf_counter()
            runTime = end - start
            # print("运行时间：", runTime * 1000)

    def _realTimeControl(self):
        result = self._controlfun(actualQ=self._actualQ, actualQd=self._actualQd.copy(),
                                                            targetQd=self._targetQd, targetQdd=self._targetQdd.copy(),
                                                            actualTcpPose=self._actualTcpPose.copy(),
                                                            actualTcpSpeed=self._actualTcpSpeed.copy(),
                                                            force=self._ft.copy())#输入当前状态，反馈目标伺服参数

        self._controlResult=result['ControlResult']#解析结果
        pose=result['ServoParm'][0]
        lookahead_time=result['ServoParm'][1]
        gain=result['ServoParm'][2]

        if (self._controlResult != ControlResult.Running):
            if (self._controlResult == ControlResult.Stop or
                self._controlResult == ControlResult.Finish):
                self.servoStop()
            self._startServoControl = False
        else:
            if (result['ControlMode'] == ControlMode.servoL):
                self.servoL(pose,0,0,self.dt,lookahead_time,gain)
            elif(result['ControlMode']==ControlMode.servoLAT):
                self.servoLAT(self._actualTcpPose,pose,0,0,self.dt,lookahead_time,gain)
            else:
                self.servoStop()
    def getRealTimeControlResult(self)->ControlResult:
        return self._controlResult

    def setStartRealTimeControl(self,isStart:bool):
        self._startServoControl=isStart

    def setControlFun(self,func):
        self._controlfun=func

    def enableLog(self):
        self.isSave=True

    def disableLog(self):
        self.isSave=False
        return np.array(self._logData)

    #endregion